package org.robot.scm.platform.socks5;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import lombok.extern.slf4j.Slf4j;
import org.robot.scm.platform.AbstractMonitor;
import org.robot.scm.platform.server.ScmHttpServer;
import org.robot.scm.pojos.config.BaseConfig;
import org.robot.scm.pojos.enums.PlatformEnum;

@Slf4j
public class Socks5Server {

    private static final int BUFFER_SIZE = 32 * 1024; // 增大缓冲区大小以提高吞吐量

    private final int port;
    private final Map<String, String> credentials;
    private final Map<SocketChannel, ClientSession> clientSessions = new HashMap<>();
    private Selector selector;
    private ServerSocketChannel serverChannel;
    private volatile boolean running = false;
    private final Map<SocketChannel, SocketChannel> remoteToClientMap = new HashMap<>();

    public Socks5Server(int port, Map<String, String> credentials) {
        this.port = port;
        this.credentials = new HashMap<>(credentials);
    }

    public void startSocks5() throws IOException {
        serverChannel = ServerSocketChannel.open();
        serverChannel.configureBlocking(false);
        serverChannel.socket().bind(new InetSocketAddress(port));

        selector = Selector.open();
        serverChannel.register(selector, SelectionKey.OP_ACCEPT);

        running = true;
        //log.info("SOCKS5代理服务器启动在端口 " + port);

        while (running) {
            try {
                if (selector.select(1000) == 0) {
                    continue;
                }

                Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();

                while (keyIterator.hasNext()) {
                    SelectionKey key = keyIterator.next();
                    keyIterator.remove();

                    if (!key.isValid()) {
                        continue;
                    }

                    try {
                        if (key.isAcceptable()) {
                            handleAccept(key);
                        } else if (key.isReadable()) {
                            handleRead(key);
                        } else if (key.isWritable()) {
                            handleWrite(key);
                        } else if (key.isConnectable()) {
                            handleConnect(key);
                        }
                    } catch (IOException e) {
                        //log.error("处理连接时发生错误: {}", e.getMessage());
                        closeConnection((SocketChannel) key.channel());
                    }
                }
            } catch (IOException e) {
                log.error("Selector IOException: {}", e.getMessage());
            } catch (Exception e) {
                log.error("Selector Exception: {}", e.getMessage());
            }
        }
    }

    private void handleAccept(SelectionKey key) throws IOException {
        ServerSocketChannel serverSocketChannel = (ServerSocketChannel) key.channel();
        SocketChannel clientChannel = serverSocketChannel.accept();
        clientChannel.configureBlocking(false);
        clientChannel.register(selector, SelectionKey.OP_READ);
        clientSessions.put(clientChannel, new ClientSession());
        // log.info("新客户端连接: {}", clientChannel.getRemoteAddress());
    }

    private void handleRead(SelectionKey key) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        ClientSession session = clientSessions.get(channel);

        if (session == null) {
            // 使用remoteToClientMap快速查找对应的客户端
            SocketChannel clientChannel = remoteToClientMap.get(channel);
            if (clientChannel != null) {
                session = clientSessions.get(clientChannel);

                // 从远程服务器读取数据并转发给客户端
                forwardData(channel, session.remoteToClientBuffer);

                if (session.remoteToClientBuffer.position() > 0) {
                    // 关注客户端通道的写事件
                    SelectionKey clientKey = clientChannel.keyFor(selector);
                    if (clientKey != null && clientKey.isValid()) {
                        clientKey.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                    }
                }
                return;
            }
            // 找不到对应的客户端会话，关闭连接
            closeConnection(channel);
            return;
        }

        // 处理客户端通道
        if (session.remoteChannel != null && session.remoteChannel.isConnected()) {
            // 从客户端读取数据并转发给远程服务器
            forwardData(channel, session.clientToRemoteBuffer);

            if (session.clientToRemoteBuffer.position() > 0) {
                // 关注远程通道的写事件
                SelectionKey remoteKey = session.remoteChannel.keyFor(selector);
                if (remoteKey != null && remoteKey.isValid()) {
                    remoteKey.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                }
            }
        } else {
            // 处理SOCKS5协议
            handleSocks5Protocol(channel, session);
        }
    }

    private void handleWrite(SelectionKey key) throws IOException {
        SocketChannel channel = (SocketChannel) key.channel();
        ClientSession session = clientSessions.get(channel);

        if (session == null) {
            // 使用remoteToClientMap快速查找对应的客户端
            SocketChannel clientChannel = remoteToClientMap.get(channel);
            if (clientChannel != null) {
                session = clientSessions.get(clientChannel);

                // 将数据写入远程服务器
                if (session.clientToRemoteBuffer.position() > 0) {
                    session.clientToRemoteBuffer.flip();
                    int written = channel.write(session.clientToRemoteBuffer);
                    session.clientToRemoteBuffer.compact();

                    if (session.clientToRemoteBuffer.position() == 0) {
                        key.interestOps(SelectionKey.OP_READ);
                    }
                }
                return;
            }
            // 找不到对应的客户端会话，关闭连接
            closeConnection(channel);
            return;
        }

        // 处理客户端通道
        if (session.remoteChannel != null && session.remoteChannel.isOpen()) {
            // 将数据写入客户端
            if (session.remoteToClientBuffer.position() > 0) {
                session.remoteToClientBuffer.flip();
                int written = channel.write(session.remoteToClientBuffer);
                session.remoteToClientBuffer.compact();

                if (session.remoteToClientBuffer.position() == 0) {
                    key.interestOps(SelectionKey.OP_READ);
                }
            }
        }
    }

    private void handleConnect(SelectionKey key) throws IOException {
        SocketChannel remoteChannel = (SocketChannel) key.channel();
        SocketChannel clientChannel = (SocketChannel) key.attachment();
        ClientSession session = clientSessions.get(clientChannel);

        if (session == null) {
            closeConnection(remoteChannel);
            return;
        }

        try {
            if (remoteChannel.finishConnect()) {
                // log.info("连接到远程服务器: {}", remoteChannel.getRemoteAddress());
                remoteChannel.register(selector, SelectionKey.OP_READ);
                sendSocks5Response(clientChannel, (byte) 0x00, session.remoteAddress);

                // 重新关注客户端通道的读写事件
                SelectionKey clientKey = clientChannel.keyFor(selector);
                if (clientKey != null && clientKey.isValid()) {
                    clientKey.interestOps(SelectionKey.OP_READ | SelectionKey.OP_WRITE);
                }

                // 将协议缓冲区中可能剩余的数据转移到转发缓冲区
                if (session.protocolBuffer.position() > 0) {
                    session.protocolBuffer.flip();
                    session.clientToRemoteBuffer.put(session.protocolBuffer);
                    session.protocolBuffer.clear();
                }
            }
        } catch (IOException e) {
            sendSocks5Response(clientChannel, (byte) 0x05, session.remoteAddress);
            closeConnection(clientChannel);
            closeChannel(session.remoteChannel);
        }
    }

    public void stop() {
        running = false;

        // 先关闭selector和serverChannel，避免在清理连接时产生异常
        if (selector != null) {
            try {
                selector.close();
            } catch (IOException e) {
                log.error("关闭selector时发生错误: {}", e.getMessage());
            }
        }
        if (serverChannel != null) {
            try {
                serverChannel.close();
            } catch (IOException e) {
                log.error("关闭serverChannel时发生错误: {}", e.getMessage());
            }
        }

        // 然后清理所有客户端连接
        Set<SocketChannel> channels = new HashSet<>(clientSessions.keySet());
        for (SocketChannel channel : channels) {
            closeConnection(channel);
        }
        clientSessions.clear();
        remoteToClientMap.clear();
    }

    public static void registry() {
        if (!AbstractMonitor.mainConfig.getEnables().contains(PlatformEnum.SOCKS5.getVal())) {
            return;
        }
        try {
            BaseConfig config = ScmHttpServer.getTargetConfig(PlatformEnum.SOCKS5);
            String noticeSrc = config.getCurrentPlatform().getNoticeSrc();
            String[] split = noticeSrc.split(",");

            // 设置用户名和密码
            Map<String, String> credentials = new HashMap<>() {{
                for (String s : split) {
                    String[] split1 = s.split(":");
                    put(split1[0], split1[1]);
                }
            }};
            String sockProxy = config.getCurrentPlatform().getSockProxy();
            int port = Integer.parseInt(sockProxy.split(":")[1]);
            // log.info("socks5 port: {}, credentials: {}", port, credentials.keySet());
            log.info("init socks5 {} ... ok", port);

            Socks5Server server = new Socks5Server(port, credentials);
            server.startSocks5();
        } catch (IOException e) {
            log.error("init socks5 error", e);
            throw new RuntimeException(e);
        }
    }

    private void handleSocks5Protocol(SocketChannel clientChannel, ClientSession session) throws IOException {
        ByteBuffer buffer = session.protocolBuffer;
        if (buffer.remaining() == 0) {
            log.warn("协议缓冲区已满，无法读取更多数据");
            return;
        }

        int bytesRead = clientChannel.read(buffer);

        if (bytesRead == -1) {
            closeConnection(clientChannel);
            return;
        }

        if (bytesRead == 0) {
            return; // 没有数据可读
        }

        buffer.flip();

        try {
            if (session.stage == ClientSession.Stage.AUTHENTICATION) {
                handleAuthentication(clientChannel, session, buffer);
            } else if (session.stage == ClientSession.Stage.REQUEST) {
                handleRequest(clientChannel, session, buffer);
            } else if (session.stage == ClientSession.Stage.CONNECTED) {
                // 连接建立后，将协议缓冲区数据转发到远程服务器
                if (buffer.hasRemaining()) {
                    session.clientToRemoteBuffer.put(buffer);
                }
            } else {
                handleInitialHandshake(clientChannel, session, buffer);
            }
        } finally {
            // 保留未处理的数据供下次解析
            if (buffer.hasRemaining()) {
                buffer.compact();
            } else {
                buffer.clear();
            }
        }
    }

    private void handleInitialHandshake(SocketChannel clientChannel, ClientSession session, ByteBuffer buffer) throws IOException {
        // 只处理完整的握手包
        if (buffer.remaining() < 2) {
            return;
        }

        buffer.mark();
        byte version = buffer.get();
        byte numMethods = buffer.get();

        // 检查是否有足够的方法数据
        if (buffer.remaining() < numMethods) {
            buffer.reset();
            return;
        }

        if (version != 0x05) {
            // 不支持的SOCKS版本
            ByteBuffer response = ByteBuffer.allocate(2);
            response.put((byte) 0x05).put((byte) 0xFF);
            response.flip();
            clientChannel.write(response);
            closeConnection(clientChannel);
            return;
        }

        // 检查是否支持用户名/密码认证
        boolean supportsAuth = false;
        byte[] methods = new byte[numMethods];
        buffer.get(methods);
        for (byte method : methods) {
            if (method == 0x02) {
                supportsAuth = true;
                break;
            }
        }

        ByteBuffer response = ByteBuffer.allocate(2);
        response.put((byte) 0x05);

        if (supportsAuth && !credentials.isEmpty()) {
            response.put((byte) 0x02); // 需要用户名/密码认证
            session.stage = ClientSession.Stage.AUTHENTICATION;
        } else if (credentials.isEmpty()) {
            response.put((byte) 0x00); // 无需认证
            session.stage = ClientSession.Stage.REQUEST;
        } else {
            response.put((byte) 0xFF); // 无 acceptable methods
            response.flip();
            clientChannel.write(response);
            closeConnection(clientChannel);
            return;
        }

        response.flip();
        clientChannel.write(response);
    }

    private void handleAuthentication(SocketChannel clientChannel, ClientSession session, ByteBuffer buffer) throws IOException {
        // 认证数据至少需要2字节(版本+用户名长度)
        if (buffer.remaining() < 2) {
            return;
        }

        buffer.mark();
        byte version = buffer.get();
        byte usernameLength = buffer.get();

        // 检查是否有足够的用户名和密码数据
        if (buffer.remaining() < usernameLength + 1) {
            buffer.reset();
            return;
        }

        if (version != 0x01) {
            // 不支持的认证版本
            ByteBuffer response = ByteBuffer.allocate(2);
            response.put((byte) 0x01).put((byte) 0xFF);
            response.flip();
            clientChannel.write(response);
            closeConnection(clientChannel);
            return;
        }

        byte[] usernameBytes = new byte[usernameLength];
        buffer.get(usernameBytes);
        String username = new String(usernameBytes);

        byte passwordLength = buffer.get();
        if (buffer.remaining() < passwordLength) {
            buffer.reset();
            return;
        }

        byte[] passwordBytes = new byte[passwordLength];
        buffer.get(passwordBytes);
        String password = new String(passwordBytes);

        // 验证凭据
        ByteBuffer response = ByteBuffer.allocate(2);
        response.put((byte) 0x01);

        if (credentials.containsKey(username) && credentials.get(username).equals(password)) {
            response.put((byte) 0x00); // 认证成功
            session.stage = ClientSession.Stage.REQUEST;
        } else {
            response.put((byte) 0x01); // 认证失败
            response.flip();
            clientChannel.write(response);
            closeConnection(clientChannel);
            return;
        }

        response.flip();
        clientChannel.write(response);
    }

    private void handleRequest(SocketChannel clientChannel, ClientSession session, ByteBuffer buffer) throws IOException {
        // 至少需要5字节基本请求数据
        if (buffer.remaining() < 5) {
            return;
        }

        buffer.mark();
        byte version = buffer.get();
        byte command = buffer.get();
        byte reserved = buffer.get();
        byte addressType = buffer.get();

        if (version != 0x05 || command != 0x01) { // 只支持CONNECT命令
            sendSocks5Response(clientChannel, (byte) 0x07, null);
            closeConnection(clientChannel);
            return;
        }

        String host = null;
        int port = 0;
        boolean hasEnoughData = true;

        switch (addressType) {
            case 0x01: // IPv4
                if (buffer.remaining() < 6) {
                    hasEnoughData = false;
                } else {
                    byte[] ipBytes = new byte[4];
                    buffer.get(ipBytes);
                    host = String.format("%d.%d.%d.%d",
                            ipBytes[0] & 0xFF, ipBytes[1] & 0xFF, ipBytes[2] & 0xFF, ipBytes[3] & 0xFF);
                    port = buffer.getShort() & 0xFFFF;
                }
                break;
            case 0x03: // 域名
                if (buffer.remaining() < 1) {
                    hasEnoughData = false;
                } else {
                    byte domainLength = buffer.get();
                    if (buffer.remaining() < domainLength + 2) {
                        hasEnoughData = false;
                    } else {
                        byte[] domainBytes = new byte[domainLength];
                        buffer.get(domainBytes);
                        host = new String(domainBytes);
                        port = buffer.getShort() & 0xFFFF;
                    }
                }
                break;
            case 0x04: // IPv6 (不支持)
                sendSocks5Response(clientChannel, (byte) 0x08, null);
                closeConnection(clientChannel);
                return;
            default:
                sendSocks5Response(clientChannel, (byte) 0x08, null);
                closeConnection(clientChannel);
                return;
        }

        if (!hasEnoughData) {
            buffer.reset();
            return;
        }

        session.remoteAddress = new InetSocketAddress(host, port);

        try {
            // 连接到目标服务器
            session.remoteChannel = SocketChannel.open();
            // 在创建远程通道时建立映射
            remoteToClientMap.put(session.remoteChannel, clientChannel);

            session.remoteChannel.configureBlocking(false);
            session.remoteChannel.connect(session.remoteAddress);
            session.remoteChannel.register(selector, SelectionKey.OP_CONNECT, clientChannel);
            session.stage = ClientSession.Stage.CONNECTED;
        } catch (IOException e) {
            sendSocks5Response(clientChannel, (byte) 0x05, session.remoteAddress);
            closeConnection(clientChannel);
        }
    }

    private void sendSocks5Response(SocketChannel clientChannel, byte status, InetSocketAddress address) throws IOException {
        ByteBuffer response = ByteBuffer.allocate(256);
        response.put((byte) 0x05); // SOCKS版本
        response.put(status); // 状态
        response.put((byte) 0x00); // 保留

        if (address == null) {
            response.put((byte) 0x01); // IPv4
            response.put(new byte[4]); // 0.0.0.0
            response.putShort((short) 0); // 端口0
        } else if (address.getAddress().getAddress().length == 4) {
            response.put((byte) 0x01); // IPv4
            response.put(address.getAddress().getAddress());
            response.putShort((short) address.getPort());
        } else {
            response.put((byte) 0x04); // IPv6
            response.put(address.getAddress().getAddress());
            response.putShort((short) address.getPort());
        }

        response.flip();
        clientChannel.write(response);
    }

    private void forwardData(SocketChannel fromChannel, ByteBuffer buffer) throws IOException {
        if (buffer.remaining() == 0) {
            // log.warn("缓冲区已满，无法读取更多数据");
            return;
        }

        int read = fromChannel.read(buffer);
        if (read == -1) {
            throw new IOException("Connection closed by peer");
        }
    }

    private void closeConnection(SocketChannel channel) {
        try {
            ClientSession session = clientSessions.remove(channel);
            if (session != null && session.remoteChannel != null) {
                remoteToClientMap.remove(session.remoteChannel);
                closeChannel(session.remoteChannel);
            }
            closeChannel(channel);
            log.info("连接关闭: {}", channel.getRemoteAddress());
        } catch (IOException e) {
            log.debug("关闭连接时发生异常: {}", e.getMessage());
        }
    }

    private void closeChannel(SocketChannel channel) throws IOException {
        if (channel != null && channel.isOpen()) {
            SelectionKey key = channel.keyFor(selector);
            if (key != null) {
                key.cancel();
            }
            channel.close();
        }
    }

    private static class ClientSession {

        enum Stage {
            INIT, AUTHENTICATION, REQUEST, CONNECTED
        }

        Stage stage = Stage.INIT;
        SocketChannel remoteChannel;
        InetSocketAddress remoteAddress;
        ByteBuffer clientToRemoteBuffer = ByteBuffer.allocateDirect(BUFFER_SIZE); // 使用直接缓冲区
        ByteBuffer remoteToClientBuffer = ByteBuffer.allocateDirect(BUFFER_SIZE); // 使用直接缓冲区
        ByteBuffer protocolBuffer = ByteBuffer.allocate(BUFFER_SIZE);
    }
}
